#!/opt/local/bin/python

import csv
import time
from numpy import arange, linspace, array
from scipy.interpolate import interp1d
from PyAJDisk import parallelFindDesiredV, parallelPowerScan, parallelFindSatPower, AJDiskAnalysis
from parseDSKfile import DSKParser
import multiprocessing as mp
import cPickle as pickle

def getResults(fname, Pin, V, port):
    """A function to return a set of simulation results for a klystron.
    This is written as a separate function in order to allow it to be called within a
    multiprocessing pool."""
    Pin = float(Pin) # Pin sometimes comes as a 0-d array, so fix it
    A = AJDiskAnalysis(DSKParsObj=DSKParser(fname), port=port)
    A.setBeamVoltage(V)
    A.klysanal.setInputPower(Pin)
    A.simulate()

    I = A.klysanal.getBeamCurrent()
    eff = A.klysanal.getSimulatedTotalEfficiency()
    inputPower = A.klysanal.getInputPower()
    outputPower = A.klysanal.getSimulatedOutputPower()

    print "A.klysanal.getBeamVoltage() = %f" % (A.klysanal.getBeamVoltage())
    print "A.klysanal.getBeamCurrent() = %f" % (A.klysanal.getBeamCurrent())
    print "A.klysanal.getInputPower() = %f" % (A.klysanal.getInputPower())
    print "A.klysanal.getSimulatedOutputPower() = %f" % (A.klysanal.getSimulatedOutputPower())

    A.close()

    return (I, eff, inputPower, outputPower)

fname = "ess-new-704moredisks.dsk"
ports = arange(25000, 25201) # ports to serve the AJDisk JVM.  Shouldn't need more than 200.
pool = mp.Pool(processes=24) # 12 threaded CPUs.  24 shows a slight improvement over 12.

# Two modulator types, so two sets of limits
mod1 = {'lowerVlim': 60e3, 'higherVlim': 80e3}
mod2 = {'lowerVlim': 80e3, 'higherVlim': 100e3}

reqPower, V, result = [], [], []
powerdata = csv.reader(open("powerlist.csv", "rU"), delimiter=",")
powerdata.next() # The first line is unnecessary column titles
for line in powerdata:
    if float(line[0]) < 5e8: # Ignore 352 MHz cavities
        continue
    reqPower.append(float(line[-1]))

start = time.time()

# Loop over 10 points in the first modulator voltage range, and find the saturated klystron power
testV = linspace(mod1['lowerVlim'], mod1['higherVlim'], num=10)
satP = []
for V in testV:
    Pin, Pout = parallelFindSatPower(filename=fname, ports=ports, beamV=V)
    satP.append( Pout )
# Use the saturated power as the dependent variable, and get an 3rd order spline interpolation func
mod1['interpfunc'] = interp1d(satP, testV, kind=3, bounds_error=True)

# Loop over 10 points in the second modulator voltage range, and find the saturated klystron power
testV = linspace(mod2['lowerVlim'], mod2['higherVlim'], num=10)
satP = []
for V in testV:
    Pin, Pout = parallelFindSatPower(filename=fname, ports=ports, beamV=V)
    satP.append( Pout )
# Use the saturated power as the dependent variable, and get an 3rd order spline interpolation func
mod2['interpfunc'] = interp1d(satP, testV, kind=3, bounds_error=True)

beamV, Pin, Pout = [], [], []
I, eff, inputPower, outputPower = [], [], [], []
# Now loop over the cavities and get the appropriate modulator voltage
# If the required saturated power is within the range of the previous interpolation, then use that,
# otherwise, use the closest possible voltage within the interpolation range, and scale the klystron's
# input power appropriately.
for P in reqPower:
    try:
        beamV.append( float( interpfunc(reqPower[cavnum]) ) )
    except ValueError:
        print "Outside interpolation bounds"
        if reqPower[cavnum] < satP[0]:
            beamV.append(lowerVlim)
        else:
            beamV.append(higherVlim)
        V = beam[-1]
        Prange = linspace(0, 200, num=24)
        outPowerRange = parallelPowerScan(Prange, fname, ports, beamV=V)
        interpPinfunc = interp1d(
                outPowerRange,
                Prange, kind=3, bounds_error=True)
        Pouttemp = reqPower[cavnum]
        try:
            Pintemp  = float(interpPinfunc(Pouttemp))
        except ValueError:
            print "Pouttemp = %f" % Pouttemp
            print "Max = %f" % (outPowerRange[-1])
            raise






for cavnum in arange(len(reqPower)):
    try:
        beamV.append( float( interpfunc(reqPower[cavnum]) ) )
        V = beamV[-1]
        Pintemp, Pouttemp = parallelFindSatPower(filename=fname, ports=ports, beamV=V)
    except ValueError:
        print "Outside interpolation bounds"
        if reqPower[cavnum] < satP[0]:
            beamV.append(lowerVlim)
        else:
            beamV.append(higherVlim)
        V = beamV[-1]
        Prange = linspace(0, 200, num=24)
        outPowerRange = parallelPowerScan(Prange, fname, ports, beamV=V)
        interpPinfunc = interp1d(
                outPowerRange,
                Prange, kind=3, bounds_error=True)
        Pouttemp = reqPower[cavnum]
        try:
            Pintemp  = float(interpPinfunc(Pouttemp))
        except ValueError:
            print "Pouttemp = %f" % Pouttemp
            print "Max = %f" % (outPowerRange[-1])
            raise

    Pin.append(Pintemp)
    Pout.append(Pouttemp)

results = []
count = 0
# Now that we have the results based on interpolations/optimisations, loop over them again,
# and calculate the actual klystron parameters using those settings
# Again, use a mp pool to distribute the work over all processors
for P in Pin:
    results.append(pool.apply_async(getResults, (fname, P, beamV[count], ports[count])))
    count += 1
for i in results:
    a = i.get()
    I.append(a[0])
    eff.append(a[1])
    inputPower.append(a[2])
    outputPower.append(a[3])

pickle.dump(beamV, open("beamvoltage.p", "wb"))
pickle.dump(eff, open("efficiency.p", "wb"))
pickle.dump(I, open("beamcurrent.p", "wb"))
pickle.dump(inputPower, open("inputPower.p", "wb"))
pickle.dump(outputPower, open("outputPower.p", "wb"))

stop = time.time()
print "Time = %f" % (stop-start)
print "Stopped"

